#include "ShadowRenderPass.h"

#include <SERender/Renderer/Renderer.h>
#include <SERender/Descriptors/FrameDescriptor.h>
#include <SEngine/Render/Descriptors/SceneDescriptor.h>
#include <SECore/Ecs/Scene.h>
#include <SEngine/Ecs/Components/ModelRenderComp.h>
#include <SEngine/Render/RenderFeature/ModelRenderFeature.h>
#include <SEngine/Render/RenderFeature/UBORenderFeature.h>
#include <SEEdit/Render/RenderFeature/DebugRenderFeature.h>
#include <SEngine/Ecs/Components/LightComp.h>
#include <SERender/Resourses/UniformBuffer.h>
#include <SEngine/Ecs/Components/TransformComp.h>
#include <SERender/Resourses/Loader/TextureLoader.h>
#include <SERender/Resourses/ShadowMap.h>
#include <SEngine/Render/RenderFeature/UBORenderFeature.h>

#include <SERender/Resourses/Loader/VertexArrayLoader.h>
#include <SERender/Resourses/Loader/ShaderLoader.h>
#include <SEngine/Render/Entity/Loader/ShadowMapLoader.h>


namespace SEngine
{


    ShadowRenderPass::ShadowRenderPass(Renderer &renderer)
        : RenderPass(renderer) 
    {
        // 线光源阴影贴图
        TextureLoader<Texture2D> texLoader;
        texLoader.SetInternalFormat(DEPTH_COMPONENT);
        texLoader.SetFormat(DEPTH_COMPONENT);
        texLoader.SetType(UNSIGNED_BYTE);
        texLoader.SetParameter(TEXTURE_MAG_FILTER, LINEAR);
        texLoader.SetParameter(TEXTURE_MIN_FILTER, LINEAR);
        texLoader.SetParameter(TEXTURE_WRAP_S, CLAMP_TO_EDGE);
        texLoader.SetParameter(TEXTURE_WRAP_T, CLAMP_TO_EDGE);
        m_directLightDepthTex = texLoader.CreateTexture(2048, 2048, nullptr);

        // 线光源阴影帧缓冲
        FrameBufferLoader fboLoader;
        fboLoader.AddAttachment(m_directLightDepthTex, Attachment::DEPTH_ATTACHMENT);
        fboLoader.SetRenderBuffer(false);
        m_directLightDepthFbo = fboLoader.CreateFrameBuffer(2048, 2048);

        // 点光源阴影贴图
        TextureLoader<TextureCube> cubeTexLoader;
        cubeTexLoader.SetInternalFormat(DEPTH_COMPONENT);
        cubeTexLoader.SetFormat(DEPTH_COMPONENT);
        cubeTexLoader.SetType(UNSIGNED_BYTE);
        cubeTexLoader.SetParameter(TEXTURE_MAG_FILTER, LINEAR);
        cubeTexLoader.SetParameter(TEXTURE_MIN_FILTER, LINEAR);
        cubeTexLoader.SetParameter(TEXTURE_WRAP_S, CLAMP_TO_EDGE);
        cubeTexLoader.SetParameter(TEXTURE_WRAP_T, CLAMP_TO_EDGE);
        cubeTexLoader.SetParameter(TEXTURE_WRAP_R, CLAMP_TO_EDGE);
        m_pointLightDepthTex = cubeTexLoader.CreateTexture(1024, 1024);
        // 点光源阴影帧缓冲
        fboLoader.AddAttachment(m_pointLightDepthTex, Attachment::DEPTH_ATTACHMENT);
        fboLoader.SetRenderBuffer(false);
        m_pointLightDepthFbo = fboLoader.CreateFrameBuffer(1024, 1024);

        // 深度贴图
        texLoader.SetInternalFormat(RGB32F);
        texLoader.SetFormat(RGB);
        texLoader.SetType(FLOAT);
        texLoader.SetParameter(TEXTURE_MIN_FILTER, NEAREST);
        texLoader.SetParameter(TEXTURE_MAG_FILTER, NEAREST);
        fboLoader.AddAttachment(texLoader.CreateTexture(800, 600, nullptr), Attachment::COLOR_ATTACHMENT0);
        fboLoader.SetRenderBuffer(true);
        m_depthFbo = fboLoader.CreateFrameBuffer(800, 600);

        // 阴影渲染对象
        VertexArrayLoader<3, 2> vaoLoader(std::make_unique<CurrentDrive>());
        vaoLoader.PushVertex({-1.0f, -1.0f, 0.0f, 0.0f, 0.0f});
        vaoLoader.PushVertex({1.0f, -1.0f, 0.0f, 1.0f, 0.0f});
        vaoLoader.PushVertex({1.0f, 1.0f, 0.0f, 1.0f, 1.0f});
        vaoLoader.PushVertex({-1.0f, 1.0f, 0.0f, 0.0f, 1.0f});
        vaoLoader.PushFace({0, 1, 2});
        vaoLoader.PushFace({0, 2, 3});
        ShaderLoader shaderLoader;
        shaderLoader.Load("D:\\WorkSpace\\serious-engine\\Engine\\Shader\\DirectLightShadow.glsl");
        m_depthShadowMaterial = std::make_shared<Material>(shaderLoader.CreateShader(EngineUniformBlockBinder()));
        m_depthShadowMaterial->GetPSO().blend = true;
        m_depthShadowMaterial->GetPSO().depth = false;
        shaderLoader.Load("D:\\WorkSpace\\serious-engine\\Engine\\Shader\\PointLightShadow.glsl");
        m_cubeDepthShadowMaterial = std::make_shared<Material>(shaderLoader.CreateShader(EngineUniformBlockBinder()));
        m_cubeDepthShadowMaterial->GetPSO().blend = true;
        m_cubeDepthShadowMaterial->GetPSO().depth = false;
        m_shadowBoard.SetMesh(std::make_unique<Mesh>(vaoLoader.CreateVertexArray()));
        m_shadowBoard.SetAPIType(DrawAPIType::DrawElements);
        m_shadowBoard.SetMode(ElemMode::TRIANGLES);
    

        // 渲染单张深度图
        {
            ShaderLoader shaderLoader;
            shaderLoader.Load("D:\\WorkSpace\\serious-engine\\Engine\\Shader\\Depth.glsl");
            m_depthMaterial = std::make_shared<Material>(shaderLoader.CreateShader(EngineUniformBlockBinder()));
            m_depthMaterial->GetPSO().depth = true;
        }

        // 渲染Cube深度图(6张)
        {
            ShaderLoader shaderLoader;
            shaderLoader.Load("D:\\WorkSpace\\serious-engine\\Engine\\Shader\\CubeDepth.glsl");
            m_cubeDepthMaterial = std::make_shared<Material>(shaderLoader.CreateShader(EngineUniformBlockBinder()));
            m_cubeDepthMaterial->GetPSO().depth = true;
        }
    }

    void ShadowRenderPass::BeginFrame()
    {
        auto & sceneDes = m_renderer.AnyDescriptor::GetDescriptor<SceneDescriptor>()->get();
        auto & frameDes = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        auto & uboFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UBORenderFeature>()->get();
        auto & modelFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<ModelRenderFeature>()->get();
        // 初始化gBuffer
        // 渲染相机角度的深度图
        m_depthFbo->Resize(frameDes.w, frameDes.h);
        m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
        m_renderer.UseFrameBuffer(*m_depthFbo, [&]() {
            m_renderer.GetDrive()->ClearColor(0.0f, 0.0f, 0.0f, 0.0f);
            m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
            sceneDes.scene->Components<ModelRenderComp>([&](Actor & actor, ModelRenderComp & modelComp) {
                if (modelComp.m_modelAssets.data) {
                    modelFeature.DrawModelWithSingleMaterial(*modelComp.m_modelAssets.data, m_depthMaterial);
                }
            });
        });
        m_depthShadowMaterial->Set("u_cameraDepthMap", Shader::Sampler{m_depthFbo->GetAttachment(Attachment::COLOR_ATTACHMENT0), TextureSlot::TEXTURE0});
        m_cubeDepthShadowMaterial->Set("u_cameraDepthMap", Shader::Sampler{m_depthFbo->GetAttachment(Attachment::COLOR_ATTACHMENT0), TextureSlot::TEXTURE0});
    }
    
    Camera CalCamera(DirectLight & light, int w, int h, float zNear, float zFar, Camera viewCamera) {
        viewCamera.SetZNear(zNear);
        viewCamera.SetZFar(zFar);
        auto inpv = glm::inverse(viewCamera.GetProjectMatrix(w, h) * viewCamera.GetViewMatrix());
        auto GetWorldPos = [inpv] (glm::vec3 vertex) {
            auto v = inpv * glm::vec4{vertex, 1.0f};
            return v / v.w;
        };
        std::vector<glm::vec4> frustum{
            GetWorldPos(glm::vec3{-1.0f, -1.0f, -1.0f}),
            GetWorldPos(glm::vec3{1.0f, -1.0f, -1.0f}),
            GetWorldPos(glm::vec3{1.0f, 1.0f, -1.0f}),
            GetWorldPos(glm::vec3{-1.0f, 1.0f, -1.0f}),
            GetWorldPos(glm::vec3{-1.0f, -1.0f, 1.0f}),
            GetWorldPos(glm::vec3{1.0f, -1.0f, 1.0f}),
            GetWorldPos(glm::vec3{1.0f, 1.0f, 1.0f}),
            GetWorldPos(glm::vec3{-1.0f, 1.0f, 1.0f}),
        };
        glm::vec4 center{0.0f};
        for (int i = 0; i < 8; ++i) {
            center += (frustum[i] / 8.0f);
        }
        float radius = -1.0f;
        for (int i = 0; i < 8; ++i) {
            radius = glm::max(radius, glm::distance(center, frustum[i]));
        }
        Camera camera;
        camera.SetUp(glm::vec3{0.0f, 1.0f, 0.0f});
        camera.SetFront(light.Direct());
        camera.SetPos(glm::vec3(center) + glm::normalize(-light.Direct()) * 2.0f * radius);
        camera.SetZNear(0.1f);
        camera.SetZFar(radius * 3);
        camera.SetSize(radius);
        camera.SetType(CameraType::Ortho);
        return camera;
    }

    void ShadowRenderPass::DrawFrame()
    {
        auto & frameDes = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        auto & sceneDes = m_renderer.AnyDescriptor::GetDescriptor<SceneDescriptor>()->get();
        auto & uboFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UBORenderFeature>()->get();
        auto & modelFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<ModelRenderFeature>()->get();
        // 渲染平行光阴影图
        m_shadowBoard.SetMaterial(m_depthShadowMaterial);
        sceneDes.scene->Components<DirectLightComp, TransformComp>([&](Actor & actor, DirectLightComp & dirLight, TransformComp & transComp) {
            auto & light = std::get<DirectLight>(dirLight.light);
            
            float zNear = frameDes.camera.GetZNear();
            float zFar = frameDes.camera.GetZFar();
            auto ZValue = [=](float i, float N, float lambda = 0.7) -> float {
                return lambda * zNear * glm::pow(zFar / zNear, i / N) + (1 - lambda) * glm::mix(zNear, zFar, i / N);
            };
            Camera cameraNear = CalCamera(light, frameDes.w, frameDes.h, ZValue(0, 3), ZValue(1, 3), frameDes.camera);
            Camera cameraMid = CalCamera(light, frameDes.w, frameDes.h, ZValue(1, 3), ZValue(2, 3), frameDes.camera);
            Camera cameraFar = CalCamera(light, frameDes.w, frameDes.h, ZValue(2, 3), ZValue(3, 3), frameDes.camera);

            m_shadowBoard.GetMaterial()->Set("u_camPos", frameDes.camera.GetPos());
            m_shadowBoard.GetMaterial()->Set("u_camFront", frameDes.camera.GetFront());
            
            m_renderer.GetDrive()->Viewport(0, 0, 2048, 2048); 
            m_renderer.UseFrameBuffer(*m_directLightDepthFbo, [&]() {
                uboFeature.UpdateCameraUBOData(cameraNear, 2048, 2048);
                m_renderer.GetDrive()->Clear(DEPTH_BUFFER_BIT);
                sceneDes.scene->Components<ModelRenderComp>([&](Actor & actor, ModelRenderComp & modelComp) {
                    if (modelComp.m_modelAssets.data) {
                        modelFeature.DrawModelWithSingleMaterial(*modelComp.m_modelAssets.data, m_depthMaterial);
                    }
                });
                uboFeature.UpdateCameraUBOData(frameDes.camera, frameDes.w, frameDes.h);
            });
            m_shadowBoard.GetMaterial()->Set("u_lightDepthMap", Shader::Sampler{m_directLightDepthTex, TextureSlot::TEXTURE1});
            m_shadowBoard.GetMaterial()->Set("u_lightPV", cameraNear.GetProjectMatrix(2048, 2048) * cameraNear.GetViewMatrix());
            m_shadowBoard.GetMaterial()->Set("u_zNear", ZValue(0, 3));
            m_shadowBoard.GetMaterial()->Set("u_zFar", ZValue(1, 3));
            m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
            m_renderer.DrawEntity(m_shadowBoard);

            m_renderer.GetDrive()->Viewport(0, 0, 2048, 2048); 
            m_renderer.UseFrameBuffer(*m_directLightDepthFbo, [&]() {
                uboFeature.UpdateCameraUBOData(cameraMid, 2048, 2048);
                m_renderer.GetDrive()->Clear(DEPTH_BUFFER_BIT);
                sceneDes.scene->Components<ModelRenderComp>([&](Actor & actor, ModelRenderComp & modelComp) {
                    if (modelComp.m_modelAssets.data) {
                        modelFeature.DrawModelWithSingleMaterial(*modelComp.m_modelAssets.data, m_depthMaterial);
                    }
                });
                uboFeature.UpdateCameraUBOData(frameDes.camera, frameDes.w, frameDes.h);
            });
            m_shadowBoard.GetMaterial()->Set("u_lightDepthMap", Shader::Sampler{m_directLightDepthTex, TextureSlot::TEXTURE1});
            m_shadowBoard.GetMaterial()->Set("u_lightPV", cameraMid.GetProjectMatrix(2048, 2048) * cameraMid.GetViewMatrix());
            m_shadowBoard.GetMaterial()->Set("u_zNear", ZValue(1, 3));
            m_shadowBoard.GetMaterial()->Set("u_zFar", ZValue(2, 3));
            m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
            m_renderer.DrawEntity(m_shadowBoard);

            m_renderer.GetDrive()->Viewport(0, 0, 2048, 2048); 
            m_renderer.UseFrameBuffer(*m_directLightDepthFbo, [&]() {
                uboFeature.UpdateCameraUBOData(cameraFar, 2048, 2048);
                m_renderer.GetDrive()->Clear(DEPTH_BUFFER_BIT);
                sceneDes.scene->Components<ModelRenderComp>([&](Actor & actor, ModelRenderComp & modelComp) {
                    if (modelComp.m_modelAssets.data) {
                        modelFeature.DrawModelWithSingleMaterial(*modelComp.m_modelAssets.data, m_depthMaterial);
                    }
                });
                uboFeature.UpdateCameraUBOData(frameDes.camera, frameDes.w, frameDes.h);
                m_shadowBoard.GetMaterial()->Set("u_lightDepthMap", Shader::Sampler{m_directLightDepthTex, TextureSlot::TEXTURE1});
                m_shadowBoard.GetMaterial()->Set("u_lightPV", cameraFar.GetProjectMatrix(2048, 2048) * cameraFar.GetViewMatrix());
                m_shadowBoard.GetMaterial()->Set("u_zNear", ZValue(2, 3));
                m_shadowBoard.GetMaterial()->Set("u_zFar", ZValue(3, 3));
            });
            m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
            m_renderer.DrawEntity(m_shadowBoard);
        });


        // 渲染点光源阴影图
        m_shadowBoard.SetMaterial(m_cubeDepthShadowMaterial);
        sceneDes.scene->Components<PointLightComp, TransformComp>([&](Actor & actor, PointLightComp & ptLight, TransformComp & transComp) {
            m_renderer.GetDrive()->Viewport(0, 0, 1024, 1024); 
            auto & light = std::get<PointLight>(ptLight.light);
            std::vector<glm::vec3> fronts{
                glm::vec3{1.0f, 0.0f, 0.0f},
                glm::vec3{-1.0f, 0.0f, 0.0f},
                glm::vec3{0.0f, 1.0f, 0.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, 0.0f, 1.0f},
                glm::vec3{0.0f, 0.0f, -1.0f},
            };
            std::vector<glm::vec3> ups{
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, 0.0f, 1.0f},
                glm::vec3{0.0f, 0.0f, -1.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},

            };
            const float zFar = 500.0f;
            for (int i = 0; i < 6; ++i) {
                Camera camera;
                camera.SetFov(90.0f);
                camera.SetPos(light.Pos());
                camera.SetZNear(0.1f);
                camera.SetZFar(zFar);
                camera.SetUp(ups[i]);
                camera.SetFront(fronts[i]);
                auto vp = camera.GetProjectMatrix(1024, 1024) * camera.GetViewMatrix();
                m_cubeDepthMaterial->Set("u_shadowMat[" + std::to_string(i) + "].matrix", vp);
            }
            m_cubeDepthMaterial->Set("u_lightPos", light.Pos());
            m_cubeDepthMaterial->Set("u_zFar", zFar);
            m_renderer.UseFrameBuffer(*m_pointLightDepthFbo, [&]() {
                m_renderer.GetDrive()->Clear(DEPTH_BUFFER_BIT);
                sceneDes.scene->Components<ModelRenderComp>([&](Actor & actor, ModelRenderComp & modelComp) {
                    if (modelComp.m_modelAssets.data) {
                        modelFeature.DrawModelWithSingleMaterial(*modelComp.m_modelAssets.data, m_cubeDepthMaterial);
                    }
                });
            });
            m_shadowBoard.GetMaterial()->Set("u_lightDepthMap", Shader::Sampler{m_pointLightDepthTex, TextureSlot::TEXTURE1});
            m_shadowBoard.GetMaterial()->Set("u_lightPos", light.Pos());
            // m_shadowBoard.GetMaterial()->Set("u_lightStrength", light.Strength());
            m_shadowBoard.GetMaterial()->Set("u_zFar", zFar);
            m_shadowBoard.GetMaterial()->Set("u_bias", 0.05f);
            m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
            m_renderer.DrawEntity(m_shadowBoard);
        });
    }   
} // namespace SEngine
